import itk
import vtk
import time

import connect2


class test():
    #
    # Test the built in ITK->VTK function
    #
    def __init__(self):
        self.itk2vtk=itk.ImageToVTKImageFilter.IF2.New()

        self.my_itk2vtk = connect2.itk2vtk()

        self.itk_reader = itk.ImageFileReader.IF2.New()

        self.fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'

        self.itk_reader.SetFileName(self.fixedImageName)
        self.itk_reader.Update()

    def connect_builtin(self):
        start = time.clock()

        self.itk2vtk.SetInput(self.itk_reader.GetOutput())
        vtk_img = self.itk2vtk.GetOutput()
        vtk_img.Update()

        end = time.clock()

        return end-start

    def connect_my(self):
        start = time.clock()

        self.my_itk2vtk.SetInput(self.itk_reader.GetOutput())
        vtk_img = self.my_itk2vtk.GetOutput()
        vtk_img.Update()

        end = time.clock()

        return end-start


    def run_test(self):
        time1 = self.connect_builtin()
        time2 = self.connect_my()

        print "Built in function: %f, My function: %f" % (time1,time2)


if __name__ == "__main__":

    this = test()
    this.run_test()